from manim import *

class AdditionSubtractionRelationship(Scene):
    def construct(self):
        # 创建数轴
        number_line = NumberLine(
            x_range=[0, 10, 1],  # 数轴范围从 0 到 10，每个刻度为 1
            length=10,           # 数轴的长度
            color=BLUE,
            include_numbers=True,  # 显示刻度数字
        )
        number_line.shift(DOWN * 2)  # 将数轴向下移动
        self.play(Create(number_line))

        # 定义整数 a 和 b
        a = 5
        b = 3

        # 显示整数 a 的单位正方形
        squares_a = self.create_unit_squares(a, color=GREEN)
        for i, square in enumerate(squares_a):
            self.play(square.animate.move_to(number_line.n2p(i + 0.5) + UP * 0.5))
        
        self.wait(1)

        # 显示整数 b 的单位正方形用于加法
        squares_b = self.create_unit_squares(b, color=ORANGE)
        for i, square in enumerate(squares_b):
            self.play(square.animate.move_to(number_line.n2p(a + i + 0.5) + UP * 0.5))
        
        self.wait(1)

        # 添加加法的文本说明
        sum_ab = a + b
        addition_text = Text(f"{a} + {b} = {sum_ab}", font_size=36).to_edge(UP)
        self.play(Write(addition_text))

        self.wait(2)

        # 移除加法正方形并显示减法的过程
        self.play(FadeOut(addition_text))
        for square in squares_b:
            self.play(FadeOut(square))
        
        self.wait(1)

        # 显示减法的过程
        subtraction_text = Text(f"{sum_ab} - {b} = {a}", font_size=36).to_edge(UP)
        self.play(Write(subtraction_text))

        # 显示整数 b 的单位正方形用于减法
        squares_b_subtraction = self.create_unit_squares(b, color=RED)
        for i, square in enumerate(squares_b_subtraction):
            square.move_to(number_line.n2p(sum_ab - i - 0.5) + UP * 0.5)
            self.add(square)

        for square in squares_b_subtraction:
            self.play(square.animate.move_to(square.get_center() + UP * 2))
            self.play(FadeOut(square))
        
        self.wait(1)

        # 保持屏幕显示几秒钟
        self.wait(3)

    def create_unit_squares(self, n, color=BLUE, opacity=1.0):
        # 创建 n 个单位正方形
        squares = VGroup()
        for i in range(n):
            square = Square(side_length=1, fill_color=color, fill_opacity=opacity)
            square.move_to(RIGHT * i)  # 设置初始位置
            squares.add(square)
        self.add(squares)
        return squares

# 执行渲染
if __name__ == "__main__":
    scene = AdditionSubtractionRelationship()
    scene.render()
